# model settings
input_size = (512, 512)
norm_cfg=dict(type='SyncBN', requires_grad=True)
model = dict(
    type='Slot_detector',
    data_preprocessor=dict(
    type='DetDataPreprocessor',
    mean=[103.530, 116.280, 123.675],
    std=[1.0, 1.0, 1.0],
    bgr_to_rgb=False,
    pad_size_divisor=32),
    backbone=dict(
        type='ResNet',
        depth=34,
        num_stages=4,
        out_indices=( 2, 3),
        frozen_stages=1,
        norm_cfg=norm_cfg,
        norm_eval=True,
        style='pytorch',
        init_cfg=dict(type='Pretrained', checkpoint='torchvision://resnet34')),
    neck=dict(
        type='PAFPN',
        in_channels=[256, 512],
        out_channels=1024,
        num_outs=2),
    bbox_head=dict(
        type='Slot_head',
        in_channels=[1024, 1024],
        feat_channels=[1024, 1024],
        num_classes=3,
        stack_nums=[2, 2],
    ),
    # model training and testing settings
    train_cfg=dict(
        assigner=dict(
            type='MaxIoUAssigner',
            pos_iou_thr=0.5,
            neg_iou_thr=0.5,
            min_pos_iou=0.,
            ignore_iof_thr=-1,
            gt_max_assign_all=False),
        sampler=dict(type='PseudoSampler'),
        smoothl1_beta=1.,
        allowed_border=-1,
        pos_weight=-1,
        neg_pos_ratio=3,
        debug=False),
    test_cfg=dict(
        nms_pre=1000,
        nms=dict(type='nms', iou_threshold=0.45),
        min_bbox_size=0,
        score_thr=0.02,
        max_per_img=200))
cudnn_benchmark = True
backend_args = None


img_norm_cfg = dict(
    mean=[123.675, 116.28, 103.53],
    std=[58.395, 57.12, 57.375],
    to_rgb=True)
train_pipeline = [
    dict(type='LoadImageFromFile', backend_args=None),
    dict(type='LoadSlotAnnotations'),
    dict(type='Slot_Resize', scale=input_size, keep_ratio=True),
    dict(type='Slot_RandomRotate',rotate_range=[-180,180],rotate_ratio=1),
    # dict(type='ToTensor',keys=['mark_points','slots']),
    dict(type='PackDetInputs',meta_keys=['img_path', 'mark_points', 'slots', 'sample_idx', 'img', 'img_shape', 'ori_shape', 'scale', 'scale_factor', 'keep_ratio'])

]
test_pipeline = [
    dict(type='LoadImageFromFile', backend_args=None),
    # dict(type='LoadSlotAnnotations'),
    dict(type='Slot_Resize', scale=input_size, keep_ratio=True),
    # dict(type='Slot_RandomRotate',rotate_range=[-180,180],rotate_ratio=1),
    # dict(type='ToTensor',keys=['mark_points','slots']),
    dict(type='PackDetInputs',meta_keys=['img_path',  'img', 'img_shape', 'ori_shape', 'scale', 'scale_factor', 'keep_ratio'])
]

dataset_type = 'Slot_dataset'
data_root = '/home/u401/slot_train/ps2.0'
train_dataloader = dict(
    batch_size=24,
    num_workers=16,
    persistent_workers=True,
    sampler=dict(type='DefaultSampler', shuffle=True),
    # batch_sampler=dict(type='AspectRatioBatchSampler'),
    dataset=dict(
        type=dataset_type,
        data_root=data_root,
        ann_file='train.txt',
        pipeline=train_pipeline,
        backend_args=backend_args))

# test_dataloader = val_dataloader

max_epochs = 80
num_last_epochs = 10
train_cfg = dict(
    type='EpochBasedTrainLoop', max_epochs=max_epochs, val_interval=5)
# val_cfg = dict(type='ValLoop')
# test_cfg = dict(type='TestLoop')

base_lr = 0.0005

optim_wrapper = dict(
    type='OptimWrapper',
    optimizer=dict(type='Adam', lr=base_lr),
    clip_grad=dict(max_norm=35, norm_type=2))

# param_scheduler = [

#     dict(
#         type='MultiStepLR',
#         begin=0,
#         end=80,
#         by_epoch=True,
#         milestones=[22, 24],
#         gamma=0.1)
# ]

param_scheduler = [
    dict(
        type='LinearLR',
        start_factor=0.00025,
        by_epoch=False,
        begin=0,
        end=2000),
    dict(
        # use cosine lr from 5 to 285 epoch
        type='CosineAnnealingLR',
        eta_min=base_lr * 0.05,
        begin=10,
        T_max=max_epochs - num_last_epochs,
        end=max_epochs - num_last_epochs,
        by_epoch=True,
        convert_to_iter_based=True),
    dict(
        # use fixed lr during last 15 epochs
        type='ConstantLR',
        by_epoch=True,
        factor=1,
        begin=max_epochs - num_last_epochs,
        end=max_epochs,
    )
]




default_hooks = dict(
    checkpoint=dict(
        type='CheckpointHook',
        by_epoch=True,
        save_last=True,
        max_keep_ckpts=10,
        interval=10))
env_cfg = dict(dist_cfg=dict(backend='nccl'))
default_scope = 'mmdet'
# default_scope=cfg.get('default_scope', 'mmengine'),
vis_backends = [
    dict(type='LocalVisBackend'),
    dict(type='TensorboardVisBackend')
]
visualizer = dict(
    type='DetLocalVisualizer', vis_backends=vis_backends, name='visualizer')
# test_dataloader = dict(
#     batch_size=8,
#     num_workers=16,
#     persistent_workers=True,
#     drop_last=False,
#     sampler=dict(type='DefaultSampler', shuffle=False),
#     dataset=dict(
#         type=dataset_type,
#         data_root=data_root,
#         ann_file='test.txt',
#         pipeline=test_pipeline,
#         backend_args=backend_args))